package com.note.es.utils;


import com.alibaba.fastjson.JSONObject;

import com.note.es.connection.ESConnection;
import com.note.es.vo.*;
import org.elasticsearch.action.search.*;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.common.unit.TimeValue;
import org.elasticsearch.search.SearchHits;
import org.elasticsearch.search.aggregations.Aggregations;
import org.elasticsearch.search.aggregations.bucket.histogram.ParsedDateHistogram;
import org.elasticsearch.search.aggregations.bucket.terms.Terms;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.util.StringUtils;

import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.io.PrintWriter;
import java.util.*;

/**
 * @ClassName: SearchUtil
 * @Description: 检索相关操作
 */
public class SearchUtil {

    private static final Logger logger = LoggerFactory.getLogger(SearchUtil.class);


    /**
     * 普通检索
     */
    public SearchResponseVO search(SearchRequestVO searchRequestVO) throws IOException {
        SearchRequest searchRequest = new SearchRequest();
        searchRequest.indices(searchRequestVO.getIndexName());
        searchRequest.types(searchRequestVO.getType());
        searchRequest.searchType(SearchType.QUERY_THEN_FETCH);
        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
        searchSourceBuilder.query(searchRequestVO.getQueryBuilder());
        searchSourceBuilder.from(searchRequestVO.getStartFrom());
        searchSourceBuilder.size(searchRequestVO.getPageSize());

        if (!StringUtils.isEmpty(searchRequestVO.getSortField())) {
            searchSourceBuilder.sort(searchRequestVO.getSortField(), searchRequestVO.getSortOrder());
        }
        if (searchRequestVO.getIncludeFields().length > 0 || searchRequestVO.getExcludeFields().length > 0) {
            searchSourceBuilder.fetchSource(searchRequestVO.getIncludeFields(), searchRequestVO.getExcludeFields());
        }
        searchRequest.source(searchSourceBuilder);
        logger.info("检索索引：{},\n检索语句：{}", Arrays.toString(searchRequest.indices()), searchRequest.source().toString());
        System.out.println("检索索引：" + Arrays.toString(searchRequest.indices()) + "\n检索语句：" + searchRequest.source().toString());
        SearchResponse searchResponse = ESConnection.getClient().search(searchRequest, RequestOptions.DEFAULT);
        SearchHits searchHits = searchResponse.getHits();
        SearchResponseVO<Map<String, Object>> searchResponseVO = new SearchResponseVO<>();
        searchResponseVO.setTotalCount(searchHits.getTotalHits());
        List<Map<String, Object>> list = new ArrayList<>();
        searchHits.iterator().forEachRemaining(hit -> {
            list.add(hit.getSourceAsMap());
        });
        searchResponseVO.setData(list);
        return searchResponseVO;
    }

    /**
     * 滚动查询
     */
    public ScrollSearchResponseVO scrollSearch(ScrollSearchRequestVO scrollSearchRequestVO) throws IOException {

        if (StringUtils.isEmpty(scrollSearchRequestVO.getScrollId())) {
            SearchRequest searchRequest = new SearchRequest();
            searchRequest.indices(scrollSearchRequestVO.getIndexName());
            searchRequest.types(scrollSearchRequestVO.getType());
//            searchRequest.searchType(SearchType.QUERY_THEN_FETCH);
            SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
            searchSourceBuilder.query(scrollSearchRequestVO.getQueryBuilder());
            searchSourceBuilder.size(scrollSearchRequestVO.getLimit());
            if (scrollSearchRequestVO.getIncludeFields().length > 0 || scrollSearchRequestVO.getExcludeFields().length > 0) {
                searchSourceBuilder.fetchSource(scrollSearchRequestVO.getIncludeFields(), scrollSearchRequestVO.getExcludeFields());
            }
            searchRequest.source(searchSourceBuilder);
            searchRequest.scroll(TimeValue.timeValueSeconds(scrollSearchRequestVO.getKeepAlive()));
            logger.info("检索索引：{},\n检索语句：{}", Arrays.toString(searchRequest.indices()), searchRequest.source().toString());
            System.out.println("检索索引：" + Arrays.toString(searchRequest.indices()) + "\n检索语句：" + searchRequest.source().toString());
            SearchResponse searchResponse = ESConnection.getClient().search(searchRequest, RequestOptions.DEFAULT);
            SearchHits searchHits = searchResponse.getHits();
            ScrollSearchResponseVO<Map<String, Object>> scrollSearchResponseVO = new ScrollSearchResponseVO<>();
            scrollSearchResponseVO.setScrollId(searchResponse.getScrollId());
            scrollSearchResponseVO.setTotalCount(searchHits.getTotalHits());
            List<Map<String, Object>> list = new ArrayList<>();
            searchHits.iterator().forEachRemaining(hit -> {
                String _id = hit.getId();
                System.out.println(_id);
                hit.getSourceAsMap().put("uuid",hit.getId());
                list.add(hit.getSourceAsMap());
            });
            scrollSearchResponseVO.setData(list);
            return scrollSearchResponseVO;
        } else {
            return notFirstScrollSearch(scrollSearchRequestVO);
        }

    }

    /**
     * 非首次滚动查询
     */
    private ScrollSearchResponseVO notFirstScrollSearch(ScrollSearchRequestVO scrollSearchRequestVO) throws IOException {
        SearchScrollRequest searchScrollRequest = new SearchScrollRequest(scrollSearchRequestVO.getScrollId());
        searchScrollRequest.scroll(TimeValue.timeValueSeconds(scrollSearchRequestVO.getKeepAlive()));
        SearchResponse searchResponse = ESConnection.getClient().scroll(searchScrollRequest, RequestOptions.DEFAULT);
        SearchHits searchHits = searchResponse.getHits();
        ScrollSearchResponseVO<Map<String, Object>> scrollSearchResponseVO = new ScrollSearchResponseVO<>();
        scrollSearchResponseVO.setScrollId(searchResponse.getScrollId());
        scrollSearchResponseVO.setTotalCount(searchHits.getTotalHits());
        List<Map<String, Object>> list = new ArrayList<>();
        searchHits.iterator().forEachRemaining(hit -> {
            list.add(hit.getSourceAsMap());
        });
        scrollSearchResponseVO.setData(list);
        return scrollSearchResponseVO;
    }


    /**
     * 清除滚动查询游标
     */
    public boolean clearScroll(String scrollId) throws IOException {
        ClearScrollRequest clearScrollRequest = new ClearScrollRequest();
        clearScrollRequest.addScrollId(scrollId);
        ClearScrollResponse clearScrollResponse = ESConnection.getClient()
                .clearScroll(clearScrollRequest, RequestOptions.DEFAULT);
        return clearScrollResponse.isSucceeded();
    }


    /**
     * 聚合
     */
    public List<Map<String, Object>> aggs(AggRequestVO aggRequestVO) throws IOException {
        SearchRequest searchRequest = new SearchRequest();
        searchRequest.indices(aggRequestVO.getIndexName());
        searchRequest.types(aggRequestVO.getType());
        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
        searchSourceBuilder.query(aggRequestVO.getQueryBuilder());
        searchSourceBuilder.aggregation(aggRequestVO.getAggregationBuilder());
        searchSourceBuilder.size(0);
        searchRequest.source(searchSourceBuilder);
        logger.info("检索索引：{},\n检索语句：{}", Arrays.toString(searchRequest.indices()), searchRequest.source().toString());
        System.out.println("检索索引：" + Arrays.toString(searchRequest.indices()) + "\n检索语句：" + searchRequest.source().toString());
        Aggregations aggregations = ESConnection.getClient().search(searchRequest, RequestOptions.DEFAULT)
                .getAggregations();
        List<Map<String, Object>> list = new ArrayList<>();
        aggResultParser(aggregations, list, null);
        return list;
    }

    /**
     * 聚合结果分析
     */
    private void aggResultParser(Aggregations aggregations, List<Map<String, Object>> list, Map<String, Object> childMap) {
        aggregations.iterator().forEachRemaining(aggregation -> {
            Object obj = aggregations.get(aggregation.getName());
            if (obj instanceof Terms) {
                Terms terms = (Terms) obj;
                terms.getBuckets().forEach(bucket -> {
                    bucketParser(list, childMap, bucket.getKeyAsString(), bucket.getDocCount(), bucket.getAggregations());
                });
            } else if (obj instanceof ParsedDateHistogram) {
                ParsedDateHistogram terms = (ParsedDateHistogram) obj;
                terms.getBuckets().forEach(bucket -> {
                    bucketParser(list, childMap, bucket.getKeyAsString(), bucket.getDocCount(), bucket.getAggregations());
                });
            }
        });
    }

    /**
     * bucket 分析
     */
    private void bucketParser(List<Map<String, Object>> list,
                              Map<String, Object> childMap,
                              String keyAsString,
                              long docCount,
                              Aggregations aggregations) {
        Map<String, Object> bucketMap;
        if (childMap == null) {
            bucketMap = new HashMap<>();
            list.add(bucketMap);
        } else {
            bucketMap = childMap;
        }
        bucketMap.put(keyAsString, docCount);
        if (!aggregations.asList().isEmpty()) {
            aggResultParser(aggregations, list, bucketMap);
        }
    }

    public void searchToJson(String indexName, int pageSize, String type) throws IOException {
        File destFile = new File(ApplicationUtil.getPropertiesValue("config.properties", "url") + type + ".txt");
        FileWriter fw = new FileWriter(destFile, true);
        PrintWriter pw = new PrintWriter(fw);
        SearchRequestVO searchRequestVO = new SearchRequestVO();
        searchRequestVO.setIndexName(new String[]{indexName});
        searchRequestVO.setType(type);
        searchRequestVO.setPageSize(pageSize);
        SearchResponseVO searchResponseVO = search(searchRequestVO);
        List<HashMap<String, Object>> list = searchResponseVO.getData();

        for (HashMap<String, Object> map : list) {
            JSONObject json = new JSONObject();
            for (String key : map.keySet()) {
                json.put(key, map.get(key));
            }
            pw.println(json.toString());
        }
        pw.flush();
        fw.flush();
        pw.close();
        fw.close();
        return;
    }

}
